import random
from collections import namedtuple
import logging
from numbskull_extend.numbskulltypes import *

    
class ConstructSubgraph:

    def __init__(self, variables, features,balance):
        self.variables = variables  #因子图变量
        self.features = features    #因子图特征
        self.balance = balance      #是否需要平衡证据变量中0和1的数目


    def construct_subgraph_for_mixture(self,evidences):
        '''
        建立既有单因子，又有双因子的子图，适用于alsa
        :param evidences: 构建因子图时所需的变量，边，特征集合
        :return: Numbskull推理所需数据（weight, variable, factor, fmap, domain_mask, edges_num, var_map）
        '''
        connected_var_set, connected_edge_set, connected_feature_set = evidences
        # 平衡化：简单地限制证据变量标签即0和1的数目相等
        if self.balance:
            label0_var = set()  #标签为0的证据变量集合
            label1_var = set()  #标签为1的证据变量集合
            new_connected_var_set = set()   #平衡化后的变量集合
            for varid in connected_var_set:
                if self.variables[varid]['is_evidence'] == True and self.variables[varid]['label'] == 1:
                    label1_var.add(varid)
                elif self.variables[varid]['is_evidence'] == True and self.variables[varid]['label'] == 0:
                    label0_var.add(varid)
                else:
                    new_connected_var_set.add(varid)
            #若标签为0的证据变量比标签为1的证据变量多，则从label0_var中随机抽取len(label1_var)个变量
            if len(label1_var) <= len(label0_var):
                sampled_label0_var = set(random.sample(list(label0_var), len(label1_var)))
                balance_connected_var_set = label1_var.union(sampled_label0_var)
            #若标签为0的证据变量比标签为1的证据变量少，则从label1_var中随机抽取len(label0_var)个变量
            elif len(label1_var) > len(label0_var):
                sampled_label1_var = set(random.sample(list(label1_var), len(label0_var)))
                balance_connected_var_set = label0_var.union(sampled_label1_var)
            new_connected_var_set = new_connected_var_set.union(balance_connected_var_set)
            new_connected_edge_set = set()  #平衡化后的边集合
            new_connected_feature_set = set()   #平衡化后的特征集合
            #选取在new_connected_var_set中新的边和特征
            for edge in connected_edge_set:
                if type(edge[1]) == int and edge[1] in new_connected_var_set:
                    new_connected_edge_set.add(edge)
                    new_connected_feature_set.add(edge[0])
                elif type(edge[1]) == tuple and (
                        edge[1][0] in new_connected_var_set or edge[1][1] in new_connected_var_set):
                    new_connected_edge_set.add(edge)
                    new_connected_feature_set.add(edge[0])
            connected_var_set = new_connected_var_set
            connected_edge_set = new_connected_edge_set
            connected_feature_set = new_connected_feature_set
        var_map = dict()   #用来记录self.variables与numbskull的variable变量的映射-(self,numbskull)
        #初始化变量
        var_num = len(connected_var_set)
        variable = np.zeros(var_num, Variable)
        variable_index = 0
        for id in connected_var_set:
            variable[variable_index]["isEvidence"] = self.variables[id]['is_evidence']  #是证据变量为True，否则为False
            variable[variable_index]["initialValue"] = self.variables[id]['label']  #变量的初始值
            variable[variable_index]["dataType"] = 0    # datatype=0表示是bool变量，为1时表示非布尔变量。
            variable[variable_index]["cardinality"] = 2
            var_map[id] = variable_index     #一一记录映射关系
            variable_index += 1
        #初始化weight,多个因子可以共享同一个weight
        weight = np.zeros(len(connected_feature_set), Weight)  #weight的数目等于feature的数目
        feature_map_weight = dict()  # 记录feature id和weight id之间的映射 [feature_id,weight_id]
        weight_index = 0
        for feature_id in connected_feature_set:
            weight[weight_index]["isFixed"] = False  #因子图学习时，权重值固定为True，否则为False
            weight[weight_index]["parameterize"] = False    #权重需要参数化，则为True，否则为False
            key = (random.sample(self.features[feature_id]['weight'].keys(),1))[0]
            weight[weight_index]["initialValue"] = self.features[feature_id]['weight'][key][0] #此处一个weight可能会有很多个weight_value,赋值为第一个
            feature_map_weight[feature_id] = weight_index   #一一记录映射关系
            weight_index += 1
        #先划分单因子双因子
        binary_feature_edge = list()    #双因子边的集合
        unary_feature_edge = list()     #单因子边的集合
        for elem in connected_edge_set:
            if self.features[elem[0]]['feature_type'] == 'unary_feature':
                unary_feature_edge.append(elem)
            elif self.features[elem[0]]['feature_type'] == 'binary_feature':
                binary_feature_edge.append(elem)
        #初始化factor,fmap,edges
        edges_num = len(unary_feature_edge) + 2*len(binary_feature_edge) # 边的数目=单因子数目+2*双因子数目
        factor = np.zeros(len(unary_feature_edge)+len(binary_feature_edge), Factor)  #因子的数目=单因子数目+双因子数目
        fmap = np.zeros(edges_num, FactorToVar)  #记录边的映射关系factor[factor_index]->fmp_index，fmap[fmp_index]->var_index
        domain_mask = np.zeros(var_num, np.bool)
        edges = list()
        unary_edge = namedtuple('unary_edge', ['index', 'factorId', 'varId'])  # 单变量因子的边
        binary_edge = namedtuple('binary_edge', ['index', 'factorId', 'varId1', 'varId2'])  # 双变量因子的边
        edge = namedtuple('edge', ['index', 'factorId', 'varId'])  # 单变量因子的边
        factor_index = 0
        fmp_index = 0
        edge_index = 0
        #初始化单变量因子
        for elem in unary_feature_edge:  # [feature_id,var_id]
            var_index = elem[1]
            factor[factor_index]["factorFunction"] = 18 #选取18号因子函数
            factor[factor_index]["weightId"] = feature_map_weight[elem[0]]  #因子对应的权重id
            factor[factor_index]["featureValue"] = 1    #因子特征值
            factor[factor_index]["arity"] = 1  # 单因子度为1
            factor[factor_index]["ftv_offset"] = fmp_index  # 偏移量每次加1
            edges.append(unary_edge(edge_index, factor_index, var_map[elem[1]]))
            fmap[fmp_index]["vid"] = var_map[elem[1]]
            fmp_index += 1
            factor_index += 1
            edge_index += 1
        #初始化双变量因子
        for elem in binary_feature_edge:  # [feature_id,（var_id1，var_id2)]
            var_index = elem[1]
            factor[factor_index]["factorFunction"] = 9  #选取9号因子函数
            factor[factor_index]["weightId"] = feature_map_weight[elem[0]]  #因子对应的权重id
            factor[factor_index]["featureValue"] = 1    #因子特征值
            factor[factor_index]["arity"] = 2  # 双因子度为2
            factor[factor_index]["ftv_offset"] = fmp_index  # 偏移量每次加1
            edges.append(binary_edge(edge_index, factor_index, var_map[elem[1][0]],var_map[elem[1][1]]))
            fmap[fmp_index]["vid"] = var_map[elem[1][0]]
            fmp_index += 1
            fmap[fmp_index]["vid"] = var_map[elem[1][1]]
            fmp_index += 1
            factor_index += 1
            edge_index += 1
        logging.info("construct subgraph for mixture succeed")
        return weight, variable, factor, fmap, domain_mask, edges_num, var_map


    def construct_subgraph_for_unaryPara(self,evidences):
        '''
           # 建立只有单因子，并且所有单因子都需要参数化的子图，适用于ER
        :param evidences:
        :return:
        '''
        connected_var_set, connected_edge_set, connected_feature_set = evidences
        # 平衡化：简单地限制证据变量标签即0和1的数目相等
        if self.balance:
            label0_var = set()  #标签为0的证据变量集合
            label1_var = set()  #标签为1的证据变量集合
            new_connected_var_set = set()   #平衡化后的变量集合
            for varid in connected_var_set:
                if self.variables[varid]['is_evidence'] == True and self.variables[varid]['label'] == 1:
                    label1_var.add(varid)
                elif self.variables[varid]['is_evidence'] == True and self.variables[varid]['label'] == 0:
                    label0_var.add(varid)
                else:
                    new_connected_var_set.add(varid)
            #若标签为0的证据变量比标签为1的证据变量多，则从label0_var中随机抽取len(label1_var)个变量
            if len(label1_var) <= len(label0_var):
                sampled_label0_var = set(random.sample(list(label0_var), len(label1_var)))
                balance_connected_var_set = label1_var.union(sampled_label0_var)
            #若标签为0的证据变量比标签为1的证据变量少，则从label1_var中随机抽取len(label0_var)个变量
            elif len(label1_var) > len(label0_var):
                sampled_label1_var = set(random.sample(list(label1_var), len(label0_var)))
                balance_connected_var_set = label0_var.union(sampled_label1_var)
            new_connected_var_set = new_connected_var_set.union(balance_connected_var_set)
            new_connected_edge_set = set()  #平衡化后的边集合
            new_connected_feature_set = set()   #平衡化后的特征集合
            #选取在new_connected_var_set中新的边和特征
            for edge in connected_edge_set:
                if type(edge[1]) == int and edge[1] in new_connected_var_set:
                    new_connected_edge_set.add(edge)
                    new_connected_feature_set.add(edge[0])
                elif type(edge[1]) == tuple and (
                        edge[1][0] in new_connected_var_set or edge[1][1] in new_connected_var_set):
                    new_connected_edge_set.add(edge)
                    new_connected_feature_set.add(edge[0])
            connected_var_set = new_connected_var_set
            connected_edge_set = new_connected_edge_set
            connected_feature_set = new_connected_feature_set
        var_map = dict()  # 用来记录self.variables与numbskull的variable变量的映射-(self,numbskull)
        # 初始化variable
        var_num = len(connected_var_set)    # 变量个数=证据变量+隐变量
        variable = np.zeros(var_num, Variable)
        variable_index= 0
        for id in connected_var_set:
            variable[variable_index]["isEvidence"] = self.variables[id]['is_evidence']  #是证据变量为True，否则为False
            variable[variable_index]["initialValue"] = self.variables[id]['label']      #变量的初始值
            variable[variable_index]["dataType"] = 0  # datatype=0表示是bool变量，为1时表示非布尔变量。
            variable[variable_index]["cardinality"] = 2
            var_map[id] = variable_index    # 一一记录映射关系
            variable_index += 1
        # 初始化weight,多个因子可以共享同一个weight
        weight = np.zeros(len(connected_feature_set), Weight)  # weight的数目等于feature的数目
        feature_map_weight = dict()  # 需要记录feature id和weight id之间的映射 [feature_id,weight_id]
        weight_index = 0
        for feature_id in connected_feature_set:
            weight[weight_index]["isFixed"] = False     #因子图学习时，权重值固定为True，否则为False
            weight[weight_index]["parameterize"] = True     #权重需要参数化，则为True，否则为False
            #实体识别任务中权重的两个参数tau和alpha
            weight[weight_index]["a"] = self.features[feature_id]['tau']
            weight[weight_index]["b"] = self.features[feature_id]['alpha']
            weight[weight_index]["initialValue"] = random.uniform(-5,5)  # 此处一个weight会有很多个weight_value，此处随机初始化一个，后面应该用不上
            feature_map_weight[feature_id] = weight_index   #一一记录映射关系
            weight_index += 1

        #初始化factor,fmap,edges
        edges_num = len(connected_edge_set)         # 边的数目
        factor = np.zeros(edges_num, Factor)        # 实体识别任务中只有单因子，所以有多少个边就有多少个因子
        fmap = np.zeros(edges_num, FactorToVar)     #记录边的映射关系factor[factor_index]->fmp_index，fmap[fmp_index]->var_index
        domain_mask = np.zeros(var_num, np.bool)
        edges = list()
        edge = namedtuple('edge', ['index', 'factorId', 'varId'])  #单变量因子的边
        factor_index = 0
        fmp_index = 0
        edge_index = 0
        for elem in connected_edge_set:  # [feature_id,var_id]
            var_index = elem[1]
            factor[factor_index]["factorFunction"] = 18     #选取9号因子函数
            factor[factor_index]["weightId"] = feature_map_weight[elem[0]]  #因子对应的权重id
            factor[factor_index]["featureValue"] = self.variables[var_index]['feature_set'][elem[0]][1] #因子特征值，在实体识别任务中为特征和变量的相似度
            factor[factor_index]["arity"] = 1  # 单因子度为1
            factor[factor_index]["ftv_offset"] = fmp_index  # 偏移量每次加1
            edges.append(edge(edge_index, factor_index, var_map[elem[1]]))
            fmap[fmp_index]["vid"] = edges[factor_index][2]
            fmap[fmp_index]["x"] = self.variables[var_index]['feature_set'][elem[0]][1]  # feature_value
            fmap[fmp_index]["theta"] = self.variables[var_index]['feature_set'][elem[0]][0]  # theta
            fmp_index += 1
            factor_index += 1
            edge_index += 1

        logging.info("construct subgraph for unaryPara succeed")
        return weight, variable, factor, fmap, domain_mask, edges_num, var_map

    def construct_subgraph_for_custom(self, var_id, evidences):
        '''
        用户自定义建子图的方法
        :param var_id:
        :param evidences:
        :return:
        '''
        pass